
#include <openssl/evp.h>
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "../iface/SSLInterface.h"
#include "AESDecrypt.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe decryption.
 * 
 * @param aeskey The AES key to use for decryption. */
AESDecrypt::AESDecrypt(const AESSymmetricKey &aeskey)
{
   const char *aes_keypass = aeskey.getKeyData();
   keypassword = (aes_keypass) ? strdup(aes_keypass) : NULL;
   input = NULL;
   output = NULL;
   if (keypassword == NULL)
   {
      errprint("AES key is NULL");
      binarykey = NULL;
      binarylength = 0;
   } else
   {
      binarykey = BaseConvert::convertBase64ToBytes(keypassword, &binarylength);
   }
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe decryption.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param aeskey The AES key to use for decryption. */
AESDecrypt::AESDecrypt(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const AESSymmetricKey &aeskey)
{
   const char *aes_keypass = aeskey.getKeyData();
   keypassword = (aes_keypass) ? strdup(aes_keypass) : NULL;
   input = &iinput;
   output = &ioutput;
   if (keypassword == NULL)
   {
      errprint("AES key is NULL");
      binarykey = NULL;
      binarylength = 0;
   } else
   {
      binarykey = BaseConvert::convertBase64ToBytes(keypassword, &binarylength);
   }
}

/** Destructor for the AES decryption class. */
AESDecrypt::~AESDecrypt()
{
   StringUtils::razeString(&keypassword);
   memset(binarykey, 0, sizeof(char)*binarylength);
   binarylength = 0;
   if (binarykey) free(binarykey);
}

/** Decrypt data block and return plaintext data. The decrypted data
 * returned must be freed by the caller.
 * 
 * @param codedblock A block of coded data to decrypt.
 * 
 * @param codedsize The length of the coded data block.
 * 
 * @param[out] plainsize The length of the plaintext block of data
 * returned. The plaintext data returned must be freed by the
 * caller. */
char *AESDecrypt::processDataBlock(const char *codedblock, apr_size_t codedsize, apr_size_t *plainsize)
{
   EVP_CIPHER_CTX ctxt;
   EVP_CIPHER_CTX_init(&ctxt);
   int bytes = strlen(keypassword);
   unsigned char *binarykey = (unsigned char *)BaseConvert::convertBase64ToBytes(keypassword, &bytes);
   EVP_CipherInit_ex(&ctxt, EVP_aes_256_cbc(), NULL, NULL, NULL, 0);
   EVP_CIPHER_CTX_set_key_length(&ctxt, (AES_BITS/8));
   unsigned char *startiv = (unsigned char *)calloc(bytes, sizeof(unsigned char));
   if (bytes > (AES_BITS / 8))
      memcpy(startiv, &(binarykey[AES_BITS / 8]), sizeof(char)*(bytes - (AES_BITS/8)));
   char *result = (char *)malloc(sizeof(char)*(codedsize));
   EVP_CipherInit_ex(&ctxt, EVP_aes_256_cbc(), NULL, binarykey, startiv, 0);
   int outlen = -1;
   int outidx = 0;
   if (!EVP_CipherUpdate(&ctxt, (unsigned char *)(&(result[outidx])),
      &outlen, (const unsigned char *)codedblock, codedsize))
   {
      errprint("failed updating cipher");
      SSLCheckError();
      goto cip_cleanup;
   }
   outidx += outlen;
   outlen = -1;
   if (!EVP_CipherFinal_ex(&ctxt, (unsigned char *)(&(result[outidx])), &outlen))
   {
      errprint("failed finalising cipher");
      SSLCheckError();
      goto cip_cleanup;
   }
   if (outidx > codedsize)
   {
      errprint("overflowed cipher (expecting at most %" APR_SIZE_T_FMT ", got %" APR_SIZE_T_FMT ")", codedsize, (apr_size_t)outidx);
      goto cip_cleanup;
   }
   if (plainsize)
      *plainsize = outidx + outlen;
   EVP_CIPHER_CTX_cleanup(&ctxt);
   free(startiv); free(binarykey);
   return result;
cip_cleanup:
   EVP_CIPHER_CTX_cleanup(&ctxt);
   free(startiv); free(binarykey); free(result);
   return NULL;
}

/** Begin the process of decrypting a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int AESDecrypt::processPipeBegin()
{
   EVP_CIPHER_CTX_init(&e_ctxt);
   int bytes = strlen(keypassword);
   unsigned char *binarykey = (unsigned char *)BaseConvert::convertBase64ToBytes(keypassword, &bytes);
   
   if (!EVP_CipherInit_ex(&e_ctxt, EVP_aes_256_cbc(), NULL, NULL, NULL, 0))
   {
      errprint("could not initialise cipher");
      free(binarykey);
      return 1;
   }
   EVP_CIPHER_CTX_set_key_length(&e_ctxt, (AES_BITS/8));
   unsigned char *startiv = (unsigned char *)calloc(bytes, sizeof(unsigned char));
   if (bytes > (AES_BITS / 8))
      memcpy(startiv, &(binarykey[AES_BITS / 8]), sizeof(char)*(bytes - (AES_BITS/8)));
   if (!EVP_CipherInit_ex(&e_ctxt, EVP_aes_256_cbc(), NULL, binarykey, startiv, 0))
   {
      errprint("could not initialise cipher with key");
      free(startiv);
      free(binarykey);
      return 1;
   }
   free(startiv);
   free(binarykey);
   return 0;
}

/** Process any outstanding encrypted data on the input pipe, and
 * place any processed plaintext data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int AESDecrypt::processPipeStep()
{
   int cip_blocksize = EVP_CIPHER_CTX_block_size(&e_ctxt);
   int input_block = input->realLength();
   char *itemp_block = (char *)malloc(sizeof(char)*input_block);
   char *otemp_block = (char *)malloc(sizeof(char)*(input_block + cip_blocksize));
   
   apr_size_t readbytes = input->readFromStream(itemp_block, input_block);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != input_block)
   {
      errprint("failed to read from input stream");
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   int outlen = -1;
   if (!EVP_CipherUpdate(&e_ctxt, (unsigned char *)otemp_block,
      &outlen, (const unsigned char *)itemp_block, input_block))
   {
      errprint("failed updating cipher");
      SSLCheckError();
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   apr_size_t writebytes = output->writeToStream(otemp_block, outlen);
   if (writebytes != outlen)
   {
      errprint("failed to write data to plaintext bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   input->removeFromStream(input_block);
   free(itemp_block);
   free(otemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int AESDecrypt::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding encrypted data in the pipe, process any
 * outstanding data in the internal buffers of the decryption process
 * and finalise the decryption process.
 * 
 * @return Zero on success, non-zero on failure (one if the cipher has
 * not been finalised and this can be run again, two of the cipher has
 * been finalised but failed). */
int AESDecrypt::processPipeEnd()
{
   processPipeStep();
   int cip_blocksize = EVP_CIPHER_CTX_block_size(&e_ctxt);
   apr_size_t codeddata_blocksize = input->realLength();
   if (codeddata_blocksize >= cip_blocksize)
   {
      errprint("too much data remaining in bucket brigade (is it still being written to?)");
      return 1;
   }
   char *itemp_block = (char *)malloc(sizeof(char)*codeddata_blocksize);
   char *otemp_block = (char *)malloc(sizeof(char)*(codeddata_blocksize + cip_blocksize));

   int outlen = -1;
   if (codeddata_blocksize > 0)
   {
      apr_size_t readbytes = input->readFromStream(itemp_block, codeddata_blocksize);
      if (readbytes != codeddata_blocksize)
      {
         errprint("failed to read data from plain bucket brigade (read %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), buffer unfinished", readbytes, codeddata_blocksize);
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      if (!EVP_CipherUpdate(&e_ctxt, (unsigned char *)otemp_block,
         &outlen, (const unsigned char *)itemp_block, codeddata_blocksize))
      {
         errprint("failed updating cipher");
         SSLCheckError();
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      apr_size_t writebytes = output->writeToStream(otemp_block, outlen);
      if (writebytes != outlen)
      {
         errprint("failed to write data to plaintext bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
         free(itemp_block);
         free(otemp_block);
         return 1;
      }
      input->removeFromStream(codeddata_blocksize);
   }
   outlen = -1;
   if (!EVP_CipherFinal_ex(&e_ctxt, (unsigned char *)(otemp_block), &outlen))
   {
      errprint("failed finalising cipher");
      free(itemp_block);
      free(otemp_block);
      SSLCheckError();
      EVP_CIPHER_CTX_cleanup(&e_ctxt);
      return 2;
   }
   apr_size_t writebytes = output->writeToStream(otemp_block, outlen);
   if (writebytes != outlen)
   {
      errprint("failed to write final data to plaintext bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
      free(itemp_block);
      free(otemp_block);
      EVP_CIPHER_CTX_cleanup(&e_ctxt);
      return 2;
   }
   free(itemp_block);
   free(otemp_block);
   EVP_CIPHER_CTX_cleanup(&e_ctxt);
   return 0;
}
